import parameter
import tensorflow as tf

"""
把预测的概率向量如[[0.9,0.05,0.05],[],...[]]转换为[1,...]
"""
def Probability_To_012(y_pred):
    result = []  # 记录需要被增加权重的loss下标
    y_pred_012 = []

    for i in range(parameter.batch_size):
        # 处理y_pred_012,找到每一个概率向量中最大的下标，加入到y_pred_012
        if y_pred[i][0] > y_pred[i][1] and y_pred[i][0] > y_pred[i][2]:
            y_pred_012.append(0)
        elif y_pred[i][1] > y_pred[i][0] and y_pred[i][1] > y_pred[i][2]:
            y_pred_012.append(1)
        elif y_pred[i][2] > y_pred[i][0] and y_pred[i][2] > y_pred[i][1]:
            y_pred_012.append(2)
        else:  # 担心有一样大的情况
            y_pred_012.append(1)
    return y_pred_012#list类型，不是tensor

def Find_FN(y_pred, y):
    result = []  # 记录需要被增加权重的loss下标
    y_pred_012 = []


    for i in range(parameter.batch_size):
        # 处理y_pred_012,找到每一个概率向量中最大的下标，加入到y_pred_012
        if y_pred[i][0] > y_pred[i][1] and y_pred[i][0] > y_pred[i][2]:
            y_pred_012.append(0)
        elif y_pred[i][1] > y_pred[i][0] and y_pred[i][1] > y_pred[i][2]:
            y_pred_012.append(1)
        elif y_pred[i][2] > y_pred[i][0] and y_pred[i][2] > y_pred[i][1]:
            y_pred_012.append(2)
        else:#担心有一样大的情况
            y_pred_012.append(1)


    for i in range(parameter.batch_size):
        # FN：正例被预测为反例,增大惩罚

        # print(y_pred_012[i])
        # print(y[i])
        # print(1)
        if y_pred_012[i] == 0 and y[i] != 0:
            result.append(parameter.FN_Cost)
        else:
            result.append(1)  # 其他情况不增加惩罚

    print("Finde_FN中的result中FN_Cost出现的次数")
    print(result.count(parameter.FN_Cost))
    return result

def sensitive_loss(loss, result):

    #loss2=loss.numpy()
    for i in range(parameter.batch_size):

        #修改值非常痛苦的办法
        part1 = loss[:i]
        part2 = loss[i + 1:]
        temp=float(loss[i]*result[i])
        val = tf.constant([temp])
        new_tensor = tf.concat([part1, val, part2], axis=0)

    return new_tensor

def sensitive_module(y_pred,y,loss):
    y_pred2 = tf.argmax(y_pred, axis=1)
    y_pred2=tf.map_fn(lambda x:1 if x!=0 else 0,y_pred2)#int64
    y_pred2 = tf.cast(y_pred2, dtype=tf.int32)#int32
    y2=tf.map_fn(lambda x:1 if x==0 else 0,y)#int32
    opt1=y_pred2+y2
    opt2=tf.map_fn(lambda x:parameter.FN_Cost if x==2 else 1,opt1)#int32
    opt2=tf.cast(opt2,dtype=tf.float32)
    loss_sensitive=loss*opt2

    return loss_sensitive

if __name__ == '__main__':
    #e=tf.constant([1, 0, 0, 0, 0, 2], shape=[6,])
    e1 = tf.constant([[0.1,0.8,0.1],[0.8,0.1,0.1],[0.1,0.8,0.1]])#预测
    e2 = tf.constant([0,0,1])#真实,(20,1,1)
    loss=[2,4,9]
    sensitive_module(e1,e2,loss)
